import time
from tqsdk import TqApi, TqAuth, TqKq
from rich.console import Console
from rich.theme import Theme

# 初始化控制台
custom_theme = Theme({
    "info": "bright_cyan",
    "warning": "bright_yellow",
    "error": "bright_red",
    "success": "bright_green",
    "trade": "bold bright_magenta"
})
console = Console(theme=custom_theme)

def main():
    try:
        # 使用常用账号密码连接账户
        api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba"))
        account = api.get_account()

        console.print("正在查询账户持仓信息...", style="info")
        while True:
            #api.wait_update()
            # 获取所有持仓
            positions = api.get_position()
            if positions:
                console.print("品种代码\t持仓数量  多仓数量  空仓数量  昨仓数量  今仓数量", style="bold cyan")
                # 遍历持仓字典
                for symbol, position in positions.items():
                    pos_long = position.pos_long
                    pos_short = position.pos_short
                    pos_long_his = position.pos_long_his
                    pos_short_his = position.pos_short_his
                    pos_long_today = position.pos_long_today
                    pos_short_today = position.pos_short_today
                    total_pos = pos_long + pos_short
                    total_his = pos_long_his + pos_short_his
                    total_today = pos_long_today + pos_short_today

                    console.print(f"{symbol}\t{total_pos}\t{pos_long}\t{pos_short}\t{total_his}\t{total_today}")
            else:
                console.print("当前账户没有持仓。", style="warning")
            break
        api.close()
    except Exception as e:
        console.print(f"查询持仓信息时出错: {e}", style="error")

if __name__ == "__main__":
    main()
